package fr.mmteam.selenium.ext;

import org.apache.commons.codec.binary.Base64;

import org.apache.commons.lang.StringUtils;

import java.util.List;

import java.util.ArrayList;

import org.openqa.selenium.htmlunit.HtmlUnitDriver;

import org.slf4j.bridge.SLF4JBridgeHandler;

import java.net.URI;

import java.lang.reflect.Constructor;
import java.util.Random;
import org.jboss.netty.handler.codec.http.HttpRequest;
import org.littleshoot.proxy.ChainProxyManager;
import org.littleshoot.proxy.DefaultHttpProxyServer;
import org.littleshoot.proxy.HttpFilter;
import org.littleshoot.proxy.HttpProxyServer;
import org.littleshoot.proxy.HttpRequestFilter;
import org.littleshoot.proxy.HttpResponseFilters;
import org.openqa.selenium.Capabilities;
import org.openqa.selenium.Proxy;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.remote.CapabilityType;
import org.openqa.selenium.remote.DesiredCapabilities;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class WebDriverExtBuilder {
    
    static {
        java.util.logging.Logger rootLogger =
            java.util.logging.LogManager.getLogManager().getLogger("");
          java.util.logging.Handler[] handlers = rootLogger.getHandlers();
          for (int i = 0; i < handlers.length; i++) {
            rootLogger.removeHandler(handlers[i]);
          }
        SLF4JBridgeHandler.install();
    }

    /**
     * Logger for this class.
     */
    private final Logger logger = LoggerFactory.getLogger(WebDriverExtBuilder.class);

    private static final int TRY_COUNT = 50;

    private static final int DEFAULT_PORT = 8183;

    private static final Random RANDOM = new Random();

    private HttpProxyServer proxyServer;

    private ProxyConfiguration proxyConfiguration;

    private Class<? extends WebDriver> driverClazz;
    
    private List<HttpRequestFilter> requestFilters = new ArrayList<HttpRequestFilter>();

    private String acceptLanguage;
    
    
    public WebDriverExtBuilder() {
    }
    
    public WebDriverExtBuilder setDriverClass(final Class<? extends WebDriver> driverClazz) {
        this.driverClazz = driverClazz;
        return this;
    }

    public WebDriverExt build() {
        if (driverClazz == null) {
            throw new IllegalArgumentException("Driver class should not be null");
        }
        DesiredCapabilities capabilities = DesiredCapabilities.htmlUnit();
        capabilities.setCapability(CapabilityType.PROXY, startHttpProxy());
        capabilities.setJavascriptEnabled(true);
        WebDriver driver;
        try {
            Constructor<? extends WebDriver> c = driverClazz.getConstructor(Capabilities.class);
            driver = c.newInstance(capabilities);
        }
        catch (Exception e) {
            throw new RuntimeException("Unable to instanciate the driver.", e);
        }
        return new WebDriverExtImpl(driver);
    }

    public WebDriverExtBuilder setProxyConfiguration(ProxyConfiguration proxyConfiguration) {
        this.proxyConfiguration = proxyConfiguration;
        return this;
    }

    public ProxyConfiguration getProxyConfiguration() {
        return proxyConfiguration;
    }    
    
    public String getAcceptLanguage() {
        return acceptLanguage;
    }    
    
    /**
     * Allow to override the language of the browser.
     *
     * @param acceptLanguage For example "en_US" or "fr_FR"
     * @return
     */
    public WebDriverExtBuilder setAcceptLanguage(String acceptLanguage) {
        this.acceptLanguage = acceptLanguage;
        return this;
    }

    private Proxy startHttpProxy() {
        for (int i = 1; i <= TRY_COUNT; i++) {
            int port = getRandomPort();
            try {
                proxyServer = createDefaultHttpProxyServer(port);
                proxyServer.start();
                return seleniumProxy(port);
            }
            catch (Exception e) {
                if (i < TRY_COUNT) {
                    logger.error("Error while starting HTTP proxy. Retry...", e);
                    continue;
                }
            }
        }
        throw new RuntimeException("Unable to start HTTP proxy after " + TRY_COUNT + " retries");
    }
    
    private DefaultHttpProxyServer createDefaultHttpProxyServer(int port) {
        ChainProxyManager cpm = null;
        if (proxyConfiguration != null) {
            cpm = new ChainProxyManager() {

                public void onCommunicationError(String hostAndPort) {
                }

                public String getChainProxy(HttpRequest httpRequest) {
                    if (proxyConfiguration.useProxy(httpRequest.getUri())) {
                        return proxyConfiguration.getProxyHost() + ":" + proxyConfiguration.getProxyPort();
                    }
                    else {
                        return null;
                    }
                }
            };
        }
        if (proxyConfiguration.getUsername() != null) {
            HttpRequestFilter proxyAuth = new HttpRequestFilter() {
                @Override
                public void filter(HttpRequest httpRequest) {
                    if (proxyConfiguration.useProxy(httpRequest.getUri())) {
                        String credential = proxyConfiguration.getUsername() + ":" + StringUtils.trimToEmpty(proxyConfiguration.getPassword());
                        Base64 base64 = new Base64();
                        String encodedCredential = new String(base64.encode(credential.getBytes()));
                        httpRequest.addHeader("Proxy-Authorization", "Basic " + encodedCredential);
                    }
                }
            };
            requestFilters.add(proxyAuth);
        }
        if (acceptLanguage != null) {
            HttpRequestFilter acceptLanguageFilter = new HttpRequestFilter() {
                @Override
                public void filter(HttpRequest httpRequest) {
                    httpRequest.removeHeader("Accept-Language");
                    httpRequest.addHeader("Accept-Language", acceptLanguage);
                }
            };
            requestFilters.add(acceptLanguageFilter);
        }
        return new DefaultHttpProxyServer(port, new HttpResponseFilters() {

            public HttpFilter getFilter(String hostAndPort) {
                return null;
            }
        }, cpm, null, new HttpRequestFilter() {

            public void filter(HttpRequest httpRequest) {
                for (HttpRequestFilter filter: requestFilters) {
                    filter.filter(httpRequest);
                }
            }
        });
    }

    private org.openqa.selenium.Proxy seleniumProxy(int port) {
        Proxy proxy = new Proxy();
        proxy.setProxyType(Proxy.ProxyType.MANUAL);
        String proxyStr = String.format("localhost:%d", port);
        proxy.setHttpProxy(proxyStr);
        proxy.setSslProxy(proxyStr);

        return proxy;
    }

    private static int getRandomPort() {
        synchronized (RANDOM) {
            return DEFAULT_PORT + RANDOM.nextInt(1000);
        }
    }

}
